import pandas as pd

# 加载 CSV 文件
df = pd.read_csv('./data/merged_file.csv')

# 定义映射关系
mapping = {
    'worried': 0,
    'happy': 1,
    'neutral': 2,
    'angry': 3,
    'surprise': 4,
    'sad': 5
}

# 根据映射关系新增 label 列
df['label'] = df['discrete'].map(mapping)

# 检查是否有未映射的值
if df['label'].isnull().any():
    print("警告：以下值在 discrete 列中没有映射到对应的标签：")
    print(df[df['label'].isnull()]['discrete'].unique())

# 保存更新后的 CSV 文件
df.to_csv('./data/merged_file_with_labels.csv', index=False)

print("新增列 label 完成，已保存到 merged_file_with_labels.csv 文件中！")
